package com.jhc.yolo;

import ai.onnxruntime.OrtException;
import com.jhc.yolo.exception.NotImplementedException;
import com.jhc.yolo.pojo.ModelConfig;
import com.jhc.yolo.pojo.Yolo;
import com.jhc.yolo.pojo.YoloV5;
import com.jhc.yolo.pojo.YoloV8;

import java.io.IOException;

public class ModelFactory {


    public Yolo getModel(ModelConfig config) throws IOException, OrtException, NotImplementedException {
        String modelName = config.getModelName();
        if (modelName.equalsIgnoreCase("yolov5")) {
            return new YoloV5(config);
        }
        else if (modelName.equalsIgnoreCase("yolov8")) {
            return new YoloV8(config);
        }
        else {
            throw new NotImplementedException();
        }

    }
}
